# 不同机械臂的dh和初始theta需要酌情修改
# 根据C点位置调整算法。
import numpy as np
import math
from scipy.spatial.transform import Rotation

def eD2oD(eP_D, eD_quat, oT_e):
    oR_e = oT_e[:3, :3]
    oP_eORG = oT_e[:3, 3]
    eR_D = Rotation.from_quat([eD_quat[1], eD_quat[2], eD_quat[3], eD_quat[0]]).as_matrix()
    oR_D = oR_e@eR_D
    oP_D = iP_2_oP(oR_e, oP_eORG, eP_D)
    return oP_D, oR_D

def iP_2_oP(oR_i, oP_iORG, iP):
    oP = oR_i@iP + oP_iORG
    return oP


def triangle_angles_A(a, b, c):

    # 计算三个角的余弦值
    cos_A = (b**2 + c**2 - a**2) / (2 * b * c)
    
    # 转换为弧度
    A = np.arccos(cos_A)
    
    return A






a_1 = 0
a_2 = 0.500
a_3 = 0
d_4 = 0.500
d_6 = 0.200


d_4_L = 0.511


# 该构型通用dh表
a_0 = 0
alpha_0 = 0
d_1 = 0


a_1 = a_1
alpha_1 = -math.pi/2
d_2 = 0

alpha_2 = 0 
a_2 = a_2
d_3 = 0

alpha_3 = -math.pi/2
a_3 = a_3 
d_4 = d_4

alpha_4 = -math.pi/2
a_4 = 0
d_5 = 0


alpha_5 = -math.pi/2
a_5 = 0
d_6 = d_6

# 需要补偿theta角初始值（ik初始姿态的theta值）
theta0_bias = np.array([0, -math.pi/2, -math.pi/2, 0, -math.pi, 0])

# 实际初始姿态和ik初始姿态的偏差
# 在算fk时候既要theta0_bias又需要<加上>theta1_bias, 计算ik时只需要<减去>theta1_bias
# 因为推导ik时候默认是按照theta0_bias姿态为初始状态的
# 而theta1_bias相当于是真实初始状态相对于ik设定初始状态已有的偏差，所以要减去
theta1_bias = np.array([0, 0, 0, 0, 0, 0])
# 最后一个坐标系在单纯的fk里面，其x轴方向为a5的方向（故在作5号关节虚拟偏移小量时需要考虑a_5朝向）
# <<<<<<<<<<<<<<<<<

# 计算FK
def dh_T(a, alpha, d, theta):
    T = np.mat([[math.cos(theta), -math.sin(theta), 0, a],
    [math.sin(theta)*math.cos(alpha), math.cos(theta)*math.cos(alpha), -math.sin(alpha), -math.sin(alpha)*d],
    [math.sin(theta)*math.sin(alpha), math.cos(theta)*math.sin(alpha), math.cos(alpha), math.cos(alpha)*d],
    [0, 0, 0, 1]])
    return T
def fk(q_real_float):
    q_bias = q_real_float + theta0_bias
    T1 = dh_T(a_0, alpha_0, d_1, q_bias[0])
    T2 = dh_T(a_1, alpha_1, d_2, q_bias[1])
    T3 = dh_T(a_2, alpha_2, d_3, q_bias[2])
    T4 = dh_T(a_3, alpha_3, d_4, q_bias[3])
    T5 = dh_T(a_4, alpha_4, d_5, q_bias[4])
    T6 = dh_T(a_5, alpha_5, d_6, q_bias[5])

    T0_6  = T1*T2*T3*T4*T5*T6

    return T0_6





# 计算IK前置变量

l_AB = a_2
l_BE = a_3
l_EC = d_4
l_BC = math.sqrt(l_BE**2 + l_EC**2)

angle_BCE = triangle_angles_A(l_BE, l_EC, l_BC)



DP_C = np.array([0, 0, -d_6])


def ik_general(q_real_float, oP_DORG, oR_D):
    # 把超限判断外置
    # IK_ANALITICAL
    # 还要判断一下各个角度的范围，再据此调整flag
    ############################

    # 求解123
    oP_C = iP_2_oP(oR_D, oP_DORG, DP_C)
    proj_oP_C = math.sqrt(oP_C[0]**2+oP_C[1]**2)


    # 分类讨论，本例中需要oP_A始终指向正向
    if proj_oP_C > 0:
        if oP_C[0]>0:
            oP_A = a_1*np.array([oP_C[0],oP_C[1],0])/proj_oP_C
        else:
            oP_A = - a_1*np.array([oP_C[0],oP_C[1],0])/proj_oP_C
    else:
        oP_A = a_1*np.array([math.cos(q_real_float[0]),math.sin(q_real_float[0]),0])
    
    oP_AC = oP_C - oP_A
    if oP_C[0]>0:
        theta_1 = np.arctan2(oP_C[1], oP_C[0])
    else:
        theta_1 = np.arctan2(-oP_C[1], -oP_C[0])

    P_AC_norm = np.linalg.norm(oP_AC)

    dotOA_AC = np.dot(np.array([1,0,0]), oP_AC)

    if dotOA_AC>=0:
        Angle_AC_Axis1 = np.arccos(oP_AC[2] / P_AC_norm) #arccos值域[0, pi]
    else:
        Angle_AC_Axis1 = - np.arccos(oP_AC[2] / P_AC_norm)
        print("dotOA_AC<0")




    if P_AC_norm > (l_AB + l_BC):
        print("too far")
        flag = 0

    else:
        flag = 1
        theta_3 = math.pi - triangle_angles_A(P_AC_norm, l_AB, l_BC) + angle_BCE
        # 求P_{AC}和Axis_1的夹角, arccos的范围是0~pi
        theta_2 = Angle_AC_Axis1-triangle_angles_A(l_BC, P_AC_norm, l_AB)

    

    # 求解456
    if flag == 1:
        # 计算矩阵^{o}R_{C0}
        oP_B = np.array([(l_AB*math.sin(theta_2)+a_1)*math.cos(theta_1),(l_AB*math.sin(theta_2)+a_1)*math.sin(theta_1), l_AB*math.cos(theta_2)])
        oP_BE = np.array([-l_BE*math.cos(theta_2+theta_3)*math.cos(theta_1), (-l_BE*math.cos(theta_2+theta_3)*math.sin(theta_1)),l_BE*math.sin(theta_2+theta_3)])
        oP_E = oP_B+oP_BE
        oP_EC = oP_C-oP_E


        oZ_C0 = oP_EC/np.linalg.norm(oP_EC)
        oY_C0 = np.array([-math.sin(theta_1), math.cos(theta_1), 0])
        oX_C0 = np.cross(oY_C0,oZ_C0)
        oR_C0 = np.column_stack((oX_C0, oY_C0, oZ_C0))

        # 计算矩阵^{C0}R_{D}
        C0R_D = oR_C0.T@oR_D
        r = C0R_D
        
        # 计算theta456
        theta_5 = np.arctan2(math.sqrt(r[2][0]**2+r[2][1]**2),r[2][2])
        if theta_5 == 0:
            theta_4 = q_real_float[3]
            theta_6 = np.arctan2(-r[0][1],r[0][0]) - theta_4
        else:
            theta_4 = np.arctan2(r[1][2]/math.sin(theta_5), r[0][2]/math.sin(theta_5))
            theta_6 = np.arctan2(r[2][1]/math.sin(theta_5), -r[2][0]/math.sin(theta_5))

        # 缩小电机4的转角，把更多旋转分配给56
        if theta_4 > np.pi/2:
            theta_4 = theta_4 - np.pi
            theta_5 = -theta_5
            theta_6 = theta_6 - np.pi
        elif theta_4 < -np.pi/2:
            theta_4 = theta_4 + np.pi
            theta_5 = -theta_5
            theta_6 = theta_6 + np.pi


        
        q_star = np.array([theta_1, theta_2, theta_3, theta_4, theta_5, theta_6]) - theta1_bias
        q_star = np.mod(q_star + np.pi, 2*np.pi) - np.pi
        return q_star
    else:
        return q_real_float




def motor_angle_clip(theta = np.zeros(6)):
    theta_clip = np.zeros(6)
    theta_clip[0] = np.clip(theta[0], -np.pi/3, np.pi/3)
    theta_clip[1] = np.clip(theta[1], -1.5, 1.5)
    theta_clip[2] = np.clip(theta[2], -2.1, 2.1)
    theta_clip[3] = np.clip(theta[3], -(2*np.pi-0.1), (2*np.pi-0.1))
    theta_clip[4] = np.clip(theta[4], -1.5, 1.5)
    theta_clip[5] = np.clip(theta[5], -(2*np.pi-0.1), (2*np.pi-0.1))

    return theta_clip





# TODO 临时添加的

l_EC_L = d_4
l_BC_L = math.sqrt(l_BE**2 + l_EC_L**2)

angle_BCE_L = triangle_angles_A(l_BE, l_EC_L, l_BC_L)

def ik_general_L(q_real_float, oP_DORG, oR_D):
    # 把超限判断外置
    # IK_ANALITICAL
    # 还要判断一下各个角度的范围，再据此调整flag
    ############################

    # 求解123
    oP_C = iP_2_oP(oR_D, oP_DORG, DP_C)
    proj_oP_C = math.sqrt(oP_C[0]**2+oP_C[1]**2)


    # 分类讨论，本例中需要oP_A始终指向正向
    if proj_oP_C > 0:
        if oP_C[0]>0:
            oP_A = a_1*np.array([oP_C[0],oP_C[1],0])/proj_oP_C
        else:
            oP_A = - a_1*np.array([oP_C[0],oP_C[1],0])/proj_oP_C
    else:
        oP_A = a_1*np.array([math.cos(q_real_float[0]),math.sin(q_real_float[0]),0])
    
    oP_AC = oP_C - oP_A
    if oP_C[0]>0:
        theta_1 = np.arctan2(oP_C[1], oP_C[0])
    else:
        theta_1 = np.arctan2(-oP_C[1], -oP_C[0])

    P_AC_norm = np.linalg.norm(oP_AC)

    dotOA_AC = np.dot(np.array([1,0,0]), oP_AC)

    if dotOA_AC>=0:
        Angle_AC_Axis1 = np.arccos(oP_AC[2] / P_AC_norm) #arccos值域[0, pi]
    else:
        Angle_AC_Axis1 = - np.arccos(oP_AC[2] / P_AC_norm)
        print("dotOA_AC<0")




    if P_AC_norm > (l_AB + l_BC_L):
        print("too far")
        flag = 0

    else:
        flag = 1
        theta_3 = math.pi - triangle_angles_A(P_AC_norm, l_AB, l_BC_L) + angle_BCE_L
        # 求P_{AC}和Axis_1的夹角, arccos的范围是0~pi
        theta_2 = Angle_AC_Axis1-triangle_angles_A(l_BC_L, P_AC_norm, l_AB)

    

    # 求解456
    if flag == 1:
        # 计算矩阵^{o}R_{C0}
        oP_B = np.array([(l_AB*math.sin(theta_2)+a_1)*math.cos(theta_1),(l_AB*math.sin(theta_2)+a_1)*math.sin(theta_1), l_AB*math.cos(theta_2)])
        oP_BE = np.array([-l_BE*math.cos(theta_2+theta_3)*math.cos(theta_1), (-l_BE*math.cos(theta_2+theta_3)*math.sin(theta_1)),l_BE*math.sin(theta_2+theta_3)])
        oP_E = oP_B+oP_BE
        oP_EC = oP_C-oP_E


        oZ_C0 = oP_EC/np.linalg.norm(oP_EC)
        oY_C0 = np.array([-math.sin(theta_1), math.cos(theta_1), 0])
        oX_C0 = np.cross(oY_C0,oZ_C0)
        oR_C0 = np.column_stack((oX_C0, oY_C0, oZ_C0))

        # 计算矩阵^{C0}R_{D}
        C0R_D = oR_C0.T@oR_D
        r = C0R_D
        
        # 计算theta456
        theta_5 = np.arctan2(math.sqrt(r[2][0]**2+r[2][1]**2),r[2][2])
        if theta_5 == 0:
            theta_4 = ik_general.last_qpos[3]
            theta_6 = np.arctan2(-r[0][1],r[0][0]) - theta_4
        else:
            theta_4 = np.arctan2(r[1][2]/math.sin(theta_5), r[0][2]/math.sin(theta_5))
            theta_6 = np.arctan2(r[2][1]/math.sin(theta_5), -r[2][0]/math.sin(theta_5))

        # 缩小电机4的转角，把更多旋转分配给56
        if theta_4 > np.pi/2:
            theta_4 = theta_4 - np.pi
            theta_5 = -theta_5
            theta_6 = theta_6 - np.pi
        elif theta_4 < -np.pi/2:
            theta_4 = theta_4 + np.pi
            theta_5 = -theta_5
            theta_6 = theta_6 + np.pi


        
        q_star = np.array([theta_1, theta_2, theta_3, theta_4, theta_5, theta_6]) - theta1_bias
        q_star = np.mod(q_star + np.pi, 2*np.pi) - np.pi
        return q_star
    else:
        return q_real_float